from config import bn
from fastapi import FastAPI
from contextlib import asynccontextmanager
from config import db
from method import Bn, Log, Account, Order, Tasks
import pandas as pd
import asyncio
import sys
import json
import websockets
import random
import requests
import h5py

if sys.platform.startswith('win'):
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

user_binance: list[dict] = []
kline_binance: dict[str, list[dict]] = {}
ticker_binance: dict[str, dict] = {}
positions_binance: dict[str, list[dict]] = {}
# ticker_data
ticker_data: list[dict] = []

# https://developers.binance.com/docs/zh-CN/derivatives/usds-margined-futures/user-data-streams/Event-Balance-and-Position-Update
# 定期续期 listenKey（60 分钟过期，建议每 30 分钟刷新）
def keep_alive_listen_key(listen_key, API_KEY):
    if listen_key:
        headers = {"X-MBX-APIKEY": API_KEY}
        response = requests.put(bn.listenKey, headers=headers, params={"listenKey": listen_key})
        print("🔁 keep-alive status:", response.status_code, response.text)

# 监听 WebSocket 推送
async def listen_account_data(listen_key: str, username: str):
    ws_url = f"{bn.WSbaseUrl}/ws/{listen_key}"
    reconnect_delay = 1  # 初始重连间隔
    try:
        async with websockets.connect(ws_url) as ws:
            reconnect_delay = 1  # 初始重连间隔
            async for msg in ws:
                data = json.loads(msg)
                index = next((i for i, item in enumerate(user_binance) if item['username'] == username), -1)
                if index < 0:
                    return
                position = []
                if data['e'] == 'ACCOUNT_UPDATE':
                    assets = data['a']
                    balance = next((d for d in assets['B'] if d['a'] == 'USDT'), {})
                    for p in assets['P']:
                        tactics = user_binance[index]['tactics']
                        tacticsItem = next((d for d in tactics if d['symbol'] == p['s'].lower()), None)
                        initialMargin = float(p['pa']) * float(p['ep']) / float(tacticsItem['leverage'])
                        if tacticsItem:
                            position.append({
                                'symbol': p['s'],
                                'positionAmt': p['pa'],
                                'entryPrice': p['ep'],
                                'breakEvenPrice': p['bep'],
                                'unRealizedProfit': p['up'],
                                'positionSide': p['ps'],
                                'initialMargin': abs(initialMargin)
                            })
                    positions_binance[username] = Account.upsert_list_fast(positions_binance[username], position)
                    user_binance[index]['balance'] = {
                        'walletBalance': balance['wb'],
                        'availableBalance': Account.available_balance(balance['wb'], positions_binance[username])
                    }
                    # 账户日志
                    Log.set_log(json.dumps(positions_binance), 'assets')
                    Log.set_log(json.dumps(user_binance), 'assets')
                elif data['e'] == 'ORDER_TRADE_UPDATE':
                    # 下单时
                    order_trade_data = data['o']
                    if abs(float(order_trade_data['q'])) > 0:
                        if order_trade_data['X'] == 'NEW' and float(order_trade_data['p']) > 0:
                            tactics = user_binance[index]['tactics']
                            tacticsItem = next((d for d in tactics if d['symbol'] == order_trade_data['s'].lower()), None)
                            initialMargin = float(order_trade_data['q']) * float(order_trade_data['p']) / float(tacticsItem['leverage'])
                            position.append({
                                'symbol': order_trade_data['s'],
                                'positionAmt': order_trade_data['q'],
                                'entryPrice': order_trade_data['p'],
                                'breakEvenPrice': 0,
                                'unRealizedProfit': 0,
                                'positionSide': order_trade_data['ps'],
                                'initialMargin': abs(initialMargin)
                            })
                        elif order_trade_data['X'] == 'CANCELED':
                            position.append({
                                'symbol': order_trade_data['s'],
                                'positionAmt': 0,
                                'entryPrice': 0,
                                'breakEvenPrice': 0,
                                'unRealizedProfit': 0,
                                'positionSide': order_trade_data['ps'],
                                'initialMargin': 0
                            }) 
                        positions_binance[username] = Account.upsert_list_fast(positions_binance[username], position)
                        Log.set_log(json.dumps(positions_binance), 'assets')
    except Exception as e:
        Log.set_log(f"listen_key error: {e}", 'error')
        await asyncio.sleep(reconnect_delay + random.uniform(0, 1))
        reconnect_delay = min(reconnect_delay * 2, 30)  # 最大延迟30秒

# 定期续期 listenKey
async def auto_keep_alive(listen_key, API_KEY):
    while True:
        await asyncio.sleep(30 * 60)  # 每30分钟
        keep_alive_listen_key(listen_key, API_KEY)

# https://developers.binance.com/docs/zh-CN/derivatives/usds-margined-futures/websocket-market-streams/Individual-Symbol-Ticker-Streams
async def get_ticker(params: str, username: str):
    tactics_str = '/'.join(f"{s.lower()}@ticker" for s in params)
    url = f"{bn.WSbaseUrl}/stream?streams={tactics_str or ''}"
    reconnect_delay = 1  # 初始重连间隔
    while True:
        try:
            async with websockets.connect(url, ping_interval=20, ping_timeout=10) as ws:
                reconnect_delay = 1
                async for msg in ws:
                    data = json.loads(msg)
                    data_info = data['data']
                    ticker_binance[data_info['s']] = {
                        'symbol': data_info['s'].lower(),
                        'lastPrice': data_info['c'],
                        'averagePrice': data_info['w'],
                        'priceChangePercent': data_info['P'], 
                    }
                    Order.set_order({
                        'user_binance': user_binance,
                        'kline_binance': kline_binance,
                        'ticker_binance': ticker_binance,
                        'positions_binance': positions_binance,
                    }, username, data_info['s'])
        except websockets.exceptions.ConnectionClosedError as e:
            Log.set_log(f"get_ticker WebSocket closed: {e} - reconnecting...", 'error')
        except Exception as e:
            Log.set_log(f"get_ticker error: {e}", 'error')
        finally:
            Log.set_log(f"get_ticker error: {e}", 'error')
            await asyncio.sleep(reconnect_delay + random.uniform(0, 1))
            reconnect_delay = min(reconnect_delay * 2, 30)

# https://developers.binance.com/docs/zh-CN/derivatives/usds-margined-futures/websocket-market-streams/Kline-Candlestick-Streams
async def get_kline(params: list[str]):
    tactics_str = '/'.join(f"{s.lower()}@kline_1h" for s in params)
    url = f"{bn.WSbaseUrl}/stream?streams={tactics_str or ''}"
    reconnect_delay = 1  # 初始重连间隔

    while True:
        try:
            async with websockets.connect(url, ping_interval=20, ping_timeout=10) as ws:
                reconnect_delay = 1
                async for msg in ws:
                    data_info = json.loads(msg)
                    data = data_info['data']
                    kline = {
                        "Open": data['k']["o"],
                        "Close": data['k']["c"],
                        "High": data['k']["h"],
                        "Low": data['k']["l"],
                        "Volume": data['k']["v"],
                        "OpenTime": data['k']["t"],
                        "CloseTime": data['k']["T"],
                        "TakerBuyBaseVolume": data['k']["V"]
                    }
                    item = kline_binance[data['s']][-1]
                    if item['OpenTime'] == kline['OpenTime']:
                        kline_binance[data['s']][-1] = kline
                    else:
                        kline_binance[data['s']].append(kline)
                        # k线日志
                        Log.set_log(json.dumps(kline_binance), 'kline')
                        # 最新价格和24H涨跌日志
                        Log.set_log(json.dumps(ticker_binance), 'ticker')
        except websockets.exceptions.ConnectionClosedError as e:
            Log.set_log(f"get_kline WebSocket closed: {e} - reconnecting...", 'error')
        except Exception as e:
            Log.set_log(f"get_kline error: {e}", 'error')
        finally:
            Log.set_log(f"get_kline error: {e}", 'error')
            await asyncio.sleep(reconnect_delay + random.uniform(0, 1))
            reconnect_delay = min(reconnect_delay * 2, 30)  # 最大延迟30秒

async def place_order(username: str):
    user = next((s for s in user_binance if s['username'] == username), None)
    if user:
        try:
            while True:
                tactics = [s['symbol'] for s in user['tactics']]
                # 获取listenKey
                my_binance = Bn.Binance(username)
                listen = my_binance.synced(bn.listenKey, {}, 'post')
                listen_key = listen["listenKey"]
                # listenKey定期续期
                listen_key_task = asyncio.create_task(auto_keep_alive(listen_key, my_binance.headers['X-MBX-APIKEY']))
                # 根据listenKey获取用户币安合约账户信息变化
                binance_task = asyncio.create_task(listen_account_data(listen_key, username))
                ticker_task = asyncio.create_task(get_ticker(tactics, username))
                await asyncio.wait(
                    [listen_key_task, binance_task, ticker_task],
                    return_when=asyncio.FIRST_COMPLETED
                )
                await asyncio.sleep(2)  # 稍作等待后重连
        except asyncio.CancelledError:
            print("place_order任务接收到取消请求，准备退出...")
            for task in [binance_task, ticker_task, listen_key_task]:
                task.cancel()
            for task in [binance_task, ticker_task, listen_key_task]:
                try:
                    await task  # 等待任务退出，必要
                except asyncio.CancelledError:
                    print(f"{task.get_coro().__name__} 已取消")
            raise
        except Exception as e:
            print(f"{e}")

# k线数据初始化
def binance_klines(data):
    # 将数据转换为DataFrame
    df = pd.DataFrame(data, columns=[
        "OpenTime", "Open", "High", "Low", "Close", "Volume",
        "CloseTime", "QuoteVolume", "Trades", "TakerBuyBaseVolume",
        "TakerBuyQuoteVolume", "Ignore"
    ])
    df = df.drop(columns=["Ignore", "Trades", "QuoteVolume", "TakerBuyQuoteVolume"])
    # 将DataFrame转换为JSON格式
    result = df.to_dict(orient='records')
    return result

def get_kline_default(symbol: str):
    # 获取K线默认值
    params = {
        'symbol': symbol.upper(),
        'interval': '1h',
        'limit': 30
    }
    response = requests.get(bn.klines, params=params)
    if response.status_code == 200:
        kline_data = response.json()
        kline_binance[symbol.upper()] = binance_klines(kline_data)

def get_assets_default(username: str):
    try:
        index = next((i for i, item in enumerate(user_binance) if item.get("username") == username), -1)
        my_binance = Bn.Binance(username)
        account_info = my_binance.synced(bn.account)
        assets = [s for s in account_info['assets'] if s['asset'] == 'USDT']
        positions_info = my_binance.synced(bn.positionRisk)
        positions = [s for s in positions_info if abs(float(s['initialMargin'])) > 0]
        user_binance[index]['balance'] = {
            'walletBalance': assets[0]['walletBalance'],
            'availableBalance': assets[0]['availableBalance']
        }
        positions_binance[username] = [{
            'symbol': p['symbol'],
            'positionAmt': p["positionAmt"],
            'entryPrice': p['entryPrice'],
            'breakEvenPrice': p['breakEvenPrice'],
            'unRealizedProfit': p['unRealizedProfit'],
            'positionSide': p['positionSide'],
            'initialMargin': p['initialMargin'], 
        } for p in positions]
        # 账户日志
        Log.set_log(json.dumps(user_binance), 'assets')
        Log.set_log(json.dumps(positions_binance), 'assets')
    except Exception as e:
        Log.set_log(f"{username} 初始化失败: {e}", 'error')

user_tasks: dict[str, asyncio.Task] = {}
async def listener_def():
    global user_tasks
    try:
        with h5py.File(db.user_data, 'r') as f:
            def print_attrs(name, obj):
                if name in f:
                    data = f[name][()]
                    data = json.loads(data.decode('utf-8')) 
                    if 'binance_key' in data:
                        user_binance.append(data)
            f.visititems(print_attrs)

        tactics: list[str] = []
        for item in user_binance:
            if item['tactics'] and len(item['tactics']) > 0:
                tacticsSymbol = [s['symbol'] for s in item['tactics']]
                tactics = list(dict.fromkeys(tactics + tacticsSymbol))
        for symbol in tactics:
            get_kline_default(symbol)
        kline_task = asyncio.create_task(get_kline(tactics))
        # 获取用户账户默认值
        for item in user_binance:
            username = item['username']
            get_assets_default(username)
        # k线日志
        Log.set_log(json.dumps(kline_binance), 'kline')
        while True:
            # 获取用户账户默认值
            for item in user_binance:
                username = item['username']
                try:
                    # 启动用户独立下单任务（防止重复）
                    if username not in user_tasks or user_tasks[username].done():
                        print(f"启动任务: {username}")
                        user_tasks[username] = asyncio.create_task(place_order(username))
                    else:
                        print(f"{username} 的任务已在运行，跳过")
                except Exception as e:
                    Log.set_log(f"{username} 初始化失败: {e}", 'error')
            # 最后统一等待所有任务并发运行
            await asyncio.sleep(60)
    except asyncio.CancelledError:
        if kline_task:
            kline_task.cancel()
        for task in user_tasks.values():
            task.cancel()
        for username, task in user_tasks.items():
            try:
                await task
            except asyncio.CancelledError:
                print(f"{username} 的任务已取消")
        raise
    except Exception as e:
        Log.set_log(f"listener_def初始化失败: {e}", 'error')

# 后端获取24小时期权价格变动，保存在h5py文件中，前端用户直接获取文件中的数据
# 减少前端每次打开页面都向币安请求一次websocket
async def binance_ticker_arr():
        global ticker_data
        url = f"{bn.WSbaseUrl}/stream?streams=!ticker@arr"
        reconnect_delay = 1  # 初始重连间隔
        while True:
            try:
                async with websockets.connect(url, ping_interval=20, ping_timeout=10) as ws:
                    reconnect_delay = 1
                    async for msg in ws:
                        data = json.loads(msg)
                        ticker_data = data['data']
            except websockets.exceptions.ConnectionClosedError as e:
                Log.set_log(f"binance_ticker_arr WebSocket closed: {e} - reconnecting...", 'error')
            except Exception as e:
                Log.set_log(f"binance_ticker_arr error: {e}", 'error')            
            finally:
                Log.set_log(f"binance_ticker_arr error: {e}", 'error')
                await asyncio.sleep(reconnect_delay + random.uniform(0, 1))
                reconnect_delay = min(reconnect_delay * 2, 30)  # 最大延迟30秒

# 生命周期管理
@asynccontextmanager
async def lifespan(app: FastAPI):
    binance_ticker = asyncio.create_task(binance_ticker_arr())
    await Tasks.task_manager.start_task('listener_def', listener_def)
    try:
        yield
    finally:
        if binance_ticker:
            binance_ticker.cancel()
        await Tasks.task_manager.stop_all()